from torch_geometric.data import HeteroData

# from torch_geometric.nn import HeteroGNN

data = HeteroData()

data['paper'].x = ...  # [num_papers, num_features_paper]
data['author'].x = ...  # [num_authors, num_features_author]
data['institution'].x = ...  # [num_institutions, num_features_institution]
data['field_of_study'].x = ...  # [num_field, num_features_field]

data['paper', 'cites', 'paper'].edge_index = ...  # [2, num_edges_cites]
data['author', 'writes', 'paper'].edge_index = ...  # [2, num_edges_writes]
data['author', 'affiliated_with', 'institution'].edge_index = ...  # [2, num_edges_affiliated]
data['paper', 'has_topic', 'field_of_study'].edge_index = ...  # [2, num_edges_topic]

data['paper', 'cites', 'paper'].edge_attr = ...  # [num_edges_cites, num_features_cites]
data['author', 'writes', 'paper'].edge_attr = ...  # [num_edges_writes, num_features_writes]
data['author', 'affiliated_with', 'institution'].edge_attr = ...  # [num_edges_affiliated, num_features_affiliated]
data['paper', 'has_topic', 'field_of_study'].edge_attr = ...  # [num_edges_topic, num_features_topic]

# model = HeteroGNN(...)
# output = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)

from torch_geometric.datasets import OGB_MAG

dataset = OGB_MAG(root='./data', preprocess='metapath2vec')
data = dataset[0]
print(data.x_dict.keys())
print(data.edge_index_dict.keys())
print("####################################")
# data.edge_index_dict[('author', 'affiliated_with', 'institution')].shape, ####
# data.edge_index_dict[('author', 'writes', 'paper')].shape,  ####
# data.edge_index_dict[('paper', 'cites', 'paper')].shape, ####
# data.edge_index_dict[('paper', 'has_topic', 'field_of_study')].shape, ####
# data.edge_index_dict[('institution', 'rev_affiliated_with', 'author')].shape,
# data.edge_index_dict[('paper', 'rev_writes', 'author')].shape,
# data.edge_index_dict[('field_of_study', 'rev_has_topic', 'paper')].shape)


print("data['paper'].x 形状 [num_papers, num_features_paper] {} ".format(data['paper'].x.shape))
print("data['author'].x  形状 [num_authors, num_features_author] {} ".format(data['author'].x.shape))
print("data['institution'].x 形状 [num_institutions, num_features_institution] {} ".format(data['institution'].x.shape))
print("data['field_of_study'].x 形状 [num_field, num_features_field] {} ".format(data['field_of_study'].x.shape))

print("data['paper', 'cites', 'paper'].edge_index 形状 [2, num_edges_cites] {} ".format(
    data['paper', 'cites', 'paper'].edge_index.shape))
print("data['author', 'writes', 'paper'].edge_index 形状  [2, num_edges_writes] {} ".format(
    data['author', 'writes', 'paper'].edge_index.shape))
print("data['author', 'affiliated_with', 'institution'].edge_index 形状[2, num_edges_affiliated] {} ".format(
    data['author', 'affiliated_with', 'institution'].edge_index.shape))
print("data['paper', 'has_topic', 'field_of_study'].edge_index 形状  [2, num_edges_topic] {} ".format(
    data['paper', 'has_topic', 'field_of_study'].edge_index.shape))

paper_node_data = data['paper']
cites_edge_data = data['paper', 'cites', 'paper']
cites_edge_data = data['paper', 'paper']
cites_edge_data = data['cites']

print(data.keys())

# data['paper'].year = ...    # Setting a new paper attribute
# del data['field_of_study']  # Deleting 'field_of_study' node type
# del data['has_topic']       # Deleting 'has_topic' edge type

#获取原数据
node_types, edge_types = data.metadata()
print(node_types)
print(edge_types)

# # 放入gpu
# data = data.to('cuda:0')
# data = data.cpu()
#
# # 查看边的情况
# print(data.has_isolated_nodes())
# print(data.has_self_loops())
# print(data.is_undirected())
#
# # 转换数据
# homogeneous_data = data.to_homogeneous()
# print(homogeneous_data)

import torch_geometric.transforms as T
# ToUndirected() 函数通过为图中的所有边添加反向边，将有向图转换为（PyG 表示形式的）无向图。
data = T.ToUndirected()(data)
# 增加自环边
data = T.AddSelfLoops()(data)
# 特征归一化
data = T.NormalizeFeatures()(data)
print(data)
